package com.atguigu.ai.operator;

import cn.hutool.core.util.IdUtil;
import com.atguigu.ai.annotation.BeanType;
import com.atguigu.ai.enums.AiTypeEnum;
import com.atguigu.ai.utils.MongoUtil;
import com.atguigu.ai.vo.QueryVo;
import com.atguigu.ai.enums.LanguageEnum;
import com.atguigu.ai.enums.SystemConstant;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
//import org.springframework.ai.vectorstore.RedisVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import reactor.core.publisher.Flux;

import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@BeanType(AiTypeEnum.OLLAMA)
public class OllamaOperator implements AiOperator {

    @Autowired
    private OllamaChatModel ollamaChatModel;

    @Resource
    // private RedisVectorStore ollamaRedisVectorStore;
    private QdrantVectorStore ollamaQdrantVectorStore;

    @Autowired
    private MongoTemplate mongoTemplate;

    @Override
    public String chat(QueryVo queryVo) {
        return null;
    }

    @Override
    public Flux<String> chatStream(String model, QueryVo queryVo) {
        // 把问题记录到mongodb
        Long chatId = queryVo.getChatId();
        if (chatId != null) {
            com.atguigu.ai.pojo.Message msg = new com.atguigu.ai.pojo.Message();
            msg.setChatId(queryVo.getChatId());
            msg.setType(0);
            msg.setContent(queryVo.getMsg());
            msg.setCreateTime(new Date());
            msg.setId(IdUtil.getSnowflake().nextId());
            this.mongoTemplate.insert(msg, MongoUtil.getMessageCollection(queryVo.getChatId()));
        }

        // 查询本地知识库
        List<Document> results = ollamaQdrantVectorStore.similaritySearch(SearchRequest
                .query(queryVo.getMsg())
                .withFilterExpression(
                        new FilterExpressionBuilder()
                                .eq("projectId", queryVo.getProjectId()) // 查询当前项目的本地知识库
                                .build())
                .withTopK(SystemConstant.TOPK) // 取前10个
                .withSimilarityThreshold(SystemConstant.SIMILARITY_THRESHOLD));

        // 把本地知识库的内容作为系统提示放入
        List<Message> msgList = results.stream().map(result ->
                new SystemMessage(result.getContent())).collect(Collectors.toList());
        // 中英文切换
        msgList.add(new SystemMessage(LanguageEnum.getMsg(queryVo.getLanguage())));
        // 加入当前用户的提问
        msgList.add(new UserMessage(queryVo.getMsg()));

        // 提交到大模型获取最终结果
        Flux<ChatResponse> responseFlux = this.ollamaChatModel.stream(new Prompt(msgList, OllamaOptions.create().withModel(model)));
        return responseFlux.map(response -> response.getResult() != null
                && response.getResult().getOutput() != null
                && response.getResult().getOutput().getContent() != null
                ? response.getResult().getOutput().getContent() : "");
    }

    @Override
    public String imageUrl(QueryVo queryVo) {
        return null;
    }

    @Override
    public String imageBase64Json(QueryVo queryVo) {
        return null;
    }

    @Override
    public String textToAudio(QueryVo queryVo) {
        return null;
    }

    @Override
    public Boolean upload(Long projectId, Long knowledgeId, String content) {
        Document document = new Document(knowledgeId.toString(), content, Map.of("projectId", projectId));
        this.ollamaQdrantVectorStore.add(List.of(document));
        return true;
    }

    @Override
    public Boolean remove(String docId) {
        return this.ollamaQdrantVectorStore.delete(List.of(docId)).get();
    }

    @Override
    public Flux<String> stream(QueryVo queryVo) {
        return null;
    }
}
